package openrouter

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"

	"resty.dev/v3"
)

type RouterClient interface {
	CreateChatCompletion(ctx context.Context, request ChatCompletionRequest) (*ChatCompletionResponse, error)
	CreateChatCompletionStream(ctx context.Context, request ChatCompletionRequest, option CreateChatCompletionStreamOption) error
	GetGeneration(ctx context.Context, requestId string) (*GenerationResponse, error)
	GetModels(ctx context.Context) (*GetModelsResponse, error)
	GetCredits(ctx context.Context) (*GetCreditsResponse, error)
	Close() error
}

type Client struct {
	config      ClientConfig
	restyClient *resty.Client
}

func NewClient(apiKey string, opts ...Option) RouterClient {
	config := DefaultConfig(apiKey)

	for _, opt := range opts {
		opt(config)
	}

	return NewClientWithConfig(*config)
}

func NewClientWithConfig(config ClientConfig) RouterClient {
	return &Client{
		config:      config,
		restyClient: resty.New(),
	}
}

func (c *Client) requestWithBody(ctx context.Context, body interface{}, apiPath APIPath) (res *resty.Response, err error) {

	res, err = c.restyClient.R().
		SetHeader("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken)).
		SetHeader("Content-Type", "application/json").
		SetContext(ctx).
		SetTimeout(c.config.Timeout).
		SetBody(body).
		Post(string(apiPath))

	if err != nil {
		return
	}

	if res.StatusCode() != 200 {
		err = errors.New(res.Status())
		return
	}

	return
}

func (c *Client) requestWithQueryParams(ctx context.Context, queryParams map[string]string, apiPath APIPath) (res *resty.Response, err error) {
	res, err = c.restyClient.R().
		SetHeader("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken)).
		SetHeader("Content-Type", "application/json").
		SetContext(ctx).
		SetTimeout(c.config.Timeout).
		SetQueryParams(queryParams).
		SetDoNotParseResponse(true).
		Get(string(apiPath))

	if err != nil {
		return
	}

	if res.StatusCode() != 200 {
		err = errors.New(res.Status())
		return
	}

	return
}

func (c *Client) requestSSE(ctx context.Context, body interface{}, apiPath APIPath) (res *resty.Response, err error) {
	res, err = c.restyClient.R().
		SetHeader("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken)).
		SetHeader("Content-Type", "application/json").
		SetContext(ctx).
		SetTimeout(c.config.Timeout).
		SetBody(body).
		SetDoNotParseResponse(true).
		Post(string(apiPath))

	if err != nil {
		return
	}

	if res.StatusCode() != 200 {
		err = errors.New(res.Status())
		return
	}

	return
}

// CreateChatCompletion 调用模型发送请求（阻塞模式）
func (c *Client) CreateChatCompletion(ctx context.Context, request ChatCompletionRequest) (resp *ChatCompletionResponse, err error) {

	if request.Stream {
		return nil, fmt.Errorf("please use CreateChatCompletionStream")
	}

	// 发起请求
	resp_, err := c.requestWithBody(ctx, request, ApiPathCompletions)
	if err != nil {
		return nil, err
	}

	// 反序列化响应结果
	respBytes, err := io.ReadAll(resp_.Body)
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(respBytes, &resp)

	return
}

// CreateChatCompletionStream 调用模型发送请求（SSE）
func (c *Client) CreateChatCompletionStream(ctx context.Context, request ChatCompletionRequest, option CreateChatCompletionStreamOption) (err error) {
	request.Stream = true

	// 发起请求
	resp_, err := c.requestSSE(ctx, request, ApiPathCompletions)
	if err != nil {
		return err
	}

	scanner := bufio.NewScanner(resp_.RawResponse.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "data: ") {
			data := strings.TrimPrefix(line, "data: ")
			if data == "[DONE]" {
				break
			}

			var sse ChatCompletionSSEResponse
			unmarshalErr := json.Unmarshal([]byte(data), &sse)
			if unmarshalErr != nil {
				return unmarshalErr
			}

			if sse.Choices[0].FinishReason == "stop" {
				if option.OnEnd != nil {
					option.OnEnd()
				}
				break
			}

			if option.OnMessage != nil {
				option.OnMessage(sse)
			}
		}
	}

	return
}

// GetGeneration 获取请求元数据
func (c *Client) GetGeneration(ctx context.Context, requestId string) (resp *GenerationResponse, err error) {
	// 发起请求
	resp_, err := c.requestWithQueryParams(ctx, map[string]string{
		"id": requestId,
	}, APIPathGeneration)
	if err != nil {
		return nil, err
	}

	// 反序列化响应结果
	respBytes, err := io.ReadAll(resp_.Body)
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(respBytes, &resp)

	return
}

// GetCredits 获取Credits使用情况
func (c *Client) GetCredits(ctx context.Context) (resp *GetCreditsResponse, err error) {

	// 发起请求
	resp_, err := c.requestWithQueryParams(ctx, nil, APIPathCredits)
	if err != nil {
		return
	}

	// 反序列化响应结果
	respBytes, err := io.ReadAll(resp_.Body)
	if err != nil {
		return
	}
	err = json.Unmarshal(respBytes, &resp)

	return
}

// GetModels 获取模型列表
func (c *Client) GetModels(ctx context.Context) (*GetModelsResponse, error) {

	client := resty.New()
	defer client.Close()

	res, err := client.R().
		SetContext(ctx).
		SetHeaders(c.config.Headers).
		SetTimeout(c.config.Timeout).
		SetResult(&GetModelsResponse{}).
		Get("https://openrouter.ai/api/v1/models")

	if err != nil {
		return nil, err
	}

	return res.Result().(*GetModelsResponse), nil
}

func (c *Client) Close() (err error) {
	err = c.restyClient.Close()
	return
}
